package com.mathlearning.auth.websocket;

import com.mathlearning.auth.util.JwtUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Component
public class WebSocketInterceptor implements HandshakeInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(WebSocketInterceptor.class);

    @Autowired
    private JwtUtil jwtUtil;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        String query = request.getURI().getQuery();
        logger.info("WebSocket握手请求: {}", query);
        
        if (query == null || !query.contains("token=")) {
            logger.warn("WebSocket连接缺少token参数");
            return false;
        }
        
        // 提取token
        String token = null;
        String[] params = query.split("&");
        for (String param : params) {
            if (param.startsWith("token=")) {
                token = param.substring(6);
                break;
            }
        }
        
        if (token == null || token.isEmpty()) {
            logger.warn("WebSocket连接token为空");
            return false;
        }
        
        try {
            // 验证token并获取用户ID
            Long userId = jwtUtil.getUserIdFromToken(token);
            if (userId == null) {
                logger.warn("WebSocket连接token无效");
                return false;
            }
            
            // 将用户ID存储到WebSocket会话属性中
            attributes.put("userId", userId);
            logger.info("WebSocket握手成功，用户ID: {}", userId);
            return true;
            
        } catch (Exception e) {
            logger.error("WebSocket握手失败: {}", e.getMessage());
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            logger.error("WebSocket握手后处理失败: {}", exception.getMessage());
        } else {
            logger.info("WebSocket握手完成");
        }
    }
}

